package com.chenfan.common.extension.xxljob;

import java.lang.reflect.Method;
import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Supplier;

import com.chenfan.common.extension.context.TraceContext;
import com.chenfan.common.extension.context.TraceContextHolder;
import com.chenfan.infra.context.TenantContextHolder;
import com.xxl.job.core.biz.model.ReturnT;
import com.xxl.job.core.executor.XxlJobExecutor;
import com.xxl.job.core.handler.IJobHandler;
import com.xxl.job.core.handler.annotation.XxlJob;
import com.xxl.job.core.handler.impl.MethodJobHandler;
import org.apache.commons.collections4.MapUtils;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.boot.context.event.ApplicationReadyEvent;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationListener;
import org.springframework.core.MethodIntrospector;
import org.springframework.core.annotation.AnnotatedElementUtils;

import lombok.extern.slf4j.Slf4j;

/**
 * XxlJob-JobHandler类增强
 * @author wulg
 * @date 2022-04-27
 **/
@Slf4j
public class XxlJobApplicationListener implements ApplicationListener<ApplicationReadyEvent> {

    /**
     * 防止二次刷新
     */
    private AtomicBoolean initialed = new AtomicBoolean(false);

    private final Supplier<List<Long>> tenantIdSupplier;

    public XxlJobApplicationListener(final Supplier<List<Long>> tenantIdSupplier) {
        this.tenantIdSupplier = tenantIdSupplier;
    }

    @Override
    public void onApplicationEvent(ApplicationReadyEvent event) {
        if (initialed.compareAndSet(false,true)) {
            this.replaceMethodJobHandler(event.getApplicationContext());
        }
    }

    /**
     * 替换MethodJobHandler
     * @param applicationContext
     */
    private void replaceMethodJobHandler(ApplicationContext applicationContext) {
        String[] beanDefinitionNames = applicationContext.getBeanNamesForType(Object.class, false, true);
        for (String beanDefinitionName : beanDefinitionNames) {
            Object bean = applicationContext.getBean(beanDefinitionName);
            Map<Method, XxlJob> annotatedMethods = MethodIntrospector.selectMethods(bean.getClass(),
                    (MethodIntrospector.MetadataLookup<XxlJob>) method -> AnnotatedElementUtils.findMergedAnnotation(method, XxlJob.class));
            if(MapUtils.isEmpty(annotatedMethods)) {
                continue;
            }
            for (Map.Entry<Method, XxlJob> methodXxlJobEntry : annotatedMethods.entrySet()) {
                Method method = methodXxlJobEntry.getKey();
                XxlJob xxlJob = methodXxlJobEntry.getValue();
                Optional<IJobHandler> optional = Optional.ofNullable(XxlJobExecutor.loadJobHandler(xxlJob.value()));
                if(optional.isPresent()) {
                    MethodJobHandler jobHandler = (MethodJobHandler)optional.get();
                    TenantIsolation tenantIsolation = method.getAnnotation(TenantIsolation.class);
                    XxlJobExecutor.registJobHandler(xxlJob.value(), new IJobHandler() {
                        @Override
                        public ReturnT<String> execute(String param) throws Exception {
                            List<Long> tenantIdList = null;
                            if(Objects.nonNull(tenantIsolation)) {
                                tenantIdList = tenantIdSupplier.get();
                                log.info("【{}】获取所有系统租户ID：{}", tenantIsolation.value(), tenantIdList);
                            }
                            if(CollectionUtils.isEmpty(tenantIdList)) {
                                try {
                                    TraceContextHolder.putSceneEntry(TraceContext.SceneEntry.JOB);
                                    return jobHandler.execute(param);
                                } finally {
                                    TraceContextHolder.clear();
                                }
                            } else {
                                //遍历系统各租户管理员信息
                                for(Long tenantId : tenantIdList) {
                                    try{
                                        TraceContextHolder.putSceneEntry(TraceContext.SceneEntry.JOB);
                                        TraceContextHolder.put(TraceContext.TENANT_ID, tenantId);
                                        TenantContextHolder.setTenantId(tenantId);
                                        jobHandler.execute(param);
                                    } finally {
                                        TraceContextHolder.clear();
                                        TenantContextHolder.clear();
                                    }
                                }
                                return ReturnT.SUCCESS;
                            }
                        }
                    });
                }
            }
        }
    }
}